# -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function

import numpy as np
import pandas as pd

from .base import Feature, OpsFeature
from ._libs.rolling import rolling_slope, rolling_rsquare
from ._libs.expanding import expanding_slope, expanding_rsquare

__all__ = [
    'Ref', 'Max', 'Min', 'Sum', 'Mean', 'Std', 'Var', 'Skew', 'Kurt', \
    'Med', 'Mad', 'Slope', 'Rsquare', 'Rank', 'Quantile', 'Count', 'EMA', \
    'Add', 'Sub', 'Mul', 'Div', 'Greater', 'Less', 'Corr', 'Cov', \
    'Abs', 'Sign', 'And', 'Or', 'Not', 'Gt', 'Ge', 'Lt', 'Le', 'Eq', 'Ne', \
]

np.seterr(invalid='ignore')

#################### Element-Wise Operator ####################
class ElemOperator(OpsFeature):
    """Element-wise Operator

    Parameters
    ----------
    feature : Feature
        feature instance
    func : str
        feature operation method

    Returns
    ----------
    Feature
        feature operation output
    """
    def __init__(self, feature, func):
        self.feature = feature
        self.func = func

    def __str__(self):
        return '{}_{}'.format(type(self).__name__.lower(), self.feature)

    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        return getattr(np, self.func)(series)
    
    
class Abs(ElemOperator):
    """Feature Absolute Value
    
    Parameters
    ----------
    feature : Feature
        feature instance

    Returns
    ----------
    Feature
        a feature instance with absolute output
    """
    def __init__(self, feature):
        super(Abs, self).__init__(feature, 'abs')
        

class Sign(ElemOperator):
    """Feature Sign
    
    Parameters
    ----------
    feature : Feature
        feature instance

    Returns
    ----------
    Feature
        a feature instance with sign
    """
    def __init__(self, feature):
        super(Abs, self).__init__(feature, 'sign')


class Not(ElemOperator):
    """Not Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        feature elementwise not output
    """
    def __init__(self, feature):
        super(Not, self).__init__(feature, 'bitwise_not')

        
#################### Pair-Wise Operator ####################        
class PairOperator(OpsFeature):
    """Pair-wise operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    func : str
        operator function
    
    Returns
    ----------
    Feature:
        two features' operation output
    """
    def __init__(self, feature_left, feature_right, func):
        self.feature_left = feature_left
        self.feature_right = feature_right
        self.func = func

    def __str__(self):
        return '{}_{}_{}'.format(type(self).__name__.lower(), self.feature_left, self.feature_right)
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        assert any([isinstance(self.feature_left, Feature), self.feature_right, Feature]), \
            'at least one of two inputs is Feature instance'
        if isinstance(self.feature_left, Feature):
            series_left = self.feature_left.load(instrument, start_time, end_time, freq)
        else:
            series_left = self.feature_left # numeric value
        if isinstance(self.feature_right, Feature):
            series_right = self.feature_right.load(instrument, start_time, end_time, freq)
        else:
            series_right = self.feature_right
        return getattr(np, self.func)(series_left, series_right)
    
    
class Add(PairOperator):
    """Add Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' sum
    """
    def __init__(self, feature_left, feature_right):
        super(Add, self).__init__(feature_left, feature_right, 'add')
        

class Sub(PairOperator):
    """Subtract Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' subtraction
    """
    def __init__(self, feature_left, feature_right):
        super(Sub, self).__init__(feature_left, feature_right, 'subtract')
    

class Mul(PairOperator):
    """Multiply Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' product
    """
    def __init__(self, feature_left, feature_right):
        super(Mul, self).__init__(feature_left, feature_right, 'multiply')
    
    
class Div(PairOperator):
    """Division Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' division
    """
    def __init__(self, feature_left, feature_right):
        super(Div, self).__init__(feature_left, feature_right, 'divide')
    
    
class Greater(PairOperator):
    """Greater Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        greater elements taken from the input two features
    """
    def __init__(self, feature_left, feature_right):
        super(Greater, self).__init__(feature_left, feature_right, 'maximum')
        

class Less(PairOperator):
    """Less Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        smaller elements taken from the input two features
    """
    def __init__(self, feature_left, feature_right):
        super(Less, self).__init__(feature_left, feature_right, 'minimum')
    

class Gt(PairOperator):
    """Greater Than Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left > right`
    """
    def __init__(self, feature_left, feature_right):
        super(Gt, self).__init__(feature_left, feature_right, 'greater')
    

class Ge(PairOperator):
    """Greater Equal Than Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left >= right`
    """
    def __init__(self, feature_left, feature_right):
        super(Ge, self).__init__(feature_left, feature_right, 'greater_equal')
    

class Lt(PairOperator):
    """Less Than Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left < right`
    """
    def __init__(self, feature_left, feature_right):
        super(Lt, self).__init__(feature_left, feature_right, 'less')
    

class Le(PairOperator):
    """Less Equal Than Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left <= right`
    """
    def __init__(self, feature_left, feature_right):
        super(Le, self).__init__(feature_left, feature_right, 'less_equal')
    

class Eq(PairOperator):
    """Equal Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left == right`
    """
    def __init__(self, feature_left, feature_right):
        super(Eq, self).__init__(feature_left, feature_right, 'equal')
    

class Ne(PairOperator):
    """Not Equal Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        bool series indicate `left != right`
    """
    def __init__(self, feature_left, feature_right):
        super(Ne, self).__init__(feature_left, feature_right, 'not_equal')
    

class And(PairOperator):
    """And Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' row by row & output
    """
    def __init__(self, feature_left, feature_right):
        super(And, self).__init__(feature_left, feature_right, 'bitwise_and')


class Or(PairOperator):
    """Or Operator

    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    
    Returns
    ----------
    Feature:
        two features' row by row | outputs
    """
    def __init__(self, feature_left, feature_right):
        super(Or, self).__init__(feature_left, feature_right, 'bitwise_or')


#################### Rolling ####################    
# NOTE: methods like `rolling.mean` are optimized with cython,
# and are super faster than `rolling.apply(np.mean)`
class Rolling(OpsFeature):
    """Rolling Operator

    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size
    func : str
        rolling method

    Returns
    ----------
    Feature
        rolling outputs
    """
    def __init__(self, feature, N, func):
        self.feature = feature
        self.N = N
        self.func = func

    def __str__(self):
        return '{}_{}_{}'.format(type(self).__name__.lower(), self.feature, self.N)
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        # NOTE: remove all null check,
        # now it's user's responsibility to decide whether use features in null days
        # isnull = series.isnull() # NOTE: isnull = NaN, inf is not null
        if self.N == 0:
            series = getattr(series.expanding(min_periods=1), self.func)()
        else:
            series = getattr(series.rolling(self.N, min_periods=1), self.func)()
            # series.iloc[:self.N-1] = np.nan
        # series[isnull] = np.nan
        return series

    
class Ref(Rolling):
    """Feature Reference
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        N = 0, retrieve the first data; N > 0, retrieve data of N periods ago; N < 0, future data

    Returns
    ----------
    Feature
        a feature instance with target reference
    """
    def __init__(self, feature, N):
        super(Ref, self).__init__(feature, N, 'ref')
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        # N = 0, return first day
        if self.N == 0:
            series = pd.Series(series.iloc[0], index=series.index)
        else:
            series = series.shift(self.N) # copy
        return series        


class Mean(Rolling):
    """Rolling Mean (MA)
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling average
    """
    def __init__(self, feature, N):
        super(Mean, self).__init__(feature, N, 'mean')
        

class Sum(Rolling):
    """Rolling Sum
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling sum
    """
    def __init__(self, feature, N):
        super(Sum, self).__init__(feature, N, 'sum')
        

class Std(Rolling):
    """Rolling Std
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling std
    """
    def __init__(self, feature, N):
        super(Std, self).__init__(feature, N, 'std')


class Var(Rolling):
    """Rolling Variance
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling variance
    """
    def __init__(self, feature, N):
        super(Var, self).__init__(feature, N, 'var')


class Skew(Rolling):
    """Rolling Skewness
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling skewness
    """
    def __init__(self, feature, N):
        super(Skew, self).__init__(feature, N, 'skew')

        
class Kurt(Rolling):
    """Rolling Kurtosis
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling kurtosis
    """
    def __init__(self, feature, N):
        super(Kurt, self).__init__(feature, N, 'kurt')


class Max(Rolling):
    """Rolling Max
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling max
    """
    def __init__(self, feature, N):
        super(Max, self).__init__(feature, N, 'max')


class IdxMax(Rolling):
    """Rolling Max Index
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling max index
    """
    def __init__(self, feature, N):
        super(IdxMax, self).__init__(feature, N, 'idxmax')
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = series.expanding(min_periods=1).apply(lambda x: x.argmax()+1)
        else:
            series = series.rolling(self.N, min_periods=1).apply(lambda x: x.argmax()+1)
        return series


class Min(Rolling):
    """Rolling Min
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling min
    """
    def __init__(self, feature, N):
        super(Min, self).__init__(feature, N, 'min')


class IdxMin(Rolling):
    """Rolling Min Index
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling min index
    """
    def __init__(self, feature, N):
        super(IdxMin, self).__init__(feature, N, 'idxmin')

    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = series.expanding(min_periods=1).apply(lambda x: x.argmin()+1)
        else:
            series = series.rolling(self.N, min_periods=1).apply(lambda x: x.argmin()+1)
        return series

    
class Quantile(Rolling):
    """Rolling Quantile
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling quantile
    """
    def __init__(self, feature, N, qscore):
        super(Quantile, self).__init__(feature, N, 'quantile')
        self.qscore = qscore

    def __str__(self):
        return '{}_{}_{}_{}'.format(type(self).__name__.lower(), self.feature, self.N, self.qscore)
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = series.expanding(min_periods=1).quantile(self.qscore)
        else:
            series = series.rolling(self.N, min_periods=1).quantile(self.qscore)
        return series


class Med(Rolling):
    """Rolling Median
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling median
    """
    def __init__(self, feature, N):
        super(Med, self).__init__(feature, N, 'median')

        
class Mad(Rolling):
    """Rolling Mean Absolute Deviation
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling mean absolute deviation
    """
    def __init__(self, feature, N):
        super(Mad, self).__init__(feature, N, 'mad')

    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        # TODO: implement in Cython
        def mad(x):
            x1 = x[~np.isnan(x)]
            return np.mean(np.abs(x1 - x1.mean()))
        if self.N == 0:
            series = series.expanding(min_periods=1).apply(mad)
        else:
            series = series.rolling(self.N, min_periods=1).apply(mad)
        return series

    
class Rank(Rolling):
    """Rolling Rank (Percentile)
    
    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling rank
    """
    def __init__(self, feature, N):
        super(Rank, self).__init__(feature, N, 'rank')
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        # TODO: implement in Cython
        def rank(x):
            if np.isnan(x[-1]): return np.nan
            x1 = x[~np.isnan(x)]
            if x1.shape[0] == 0: return np.nan
            return (x1.argsort()[-1] + 1) / len(x1)
        if self.N == 0:
            series = series.expanding(min_periods=1).apply(rank)
        else:
            series = series.rolling(self.N, min_periods=1).apply(rank)
        return series


class Count(Rolling):
    """Rolling Count 

    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling count of number of non-NaN elements
    """
    def __init__(self, feature, N):
        super(Count, self).__init__(feature, N, 'count')

    
class Slope(Rolling):
    """Rolling Slope

    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with regression slope of given window
    """
    def __init__(self, feature, N):
        super(Slope, self).__init__(feature, N, 'slope')
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = pd.Series(expanding_slope(series.values), index=series.index)
        else:
            series = pd.Series(rolling_slope(series.values, self.N), index=series.index)
        return series


class Rsquare(Rolling):
    """Rolling R-value Square

    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with regression r-value square of given window
    """
    def __init__(self, feature, N):
        super(Rsquare, self).__init__(feature, N, 'rsquare')
    
    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = pd.Series(expanding_rsquare(series.values), index=series.index)
        else:
            series = pd.Series(rolling_rsquare(series.values, self.N), index=series.index)
        return series
    

class EMA(Rolling):
    """Rolling Exponential Mean (EMA)

    Parameters
    ----------
    feature : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with regression r-value square of given window
    """
    def __init__(self, feature, N):
        super(EMA, self).__init__(feature, N, 'ema')

    def _load_internal(self, instrument, start_time, end_time, freq):
        series = self.feature.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = series.ewm(len(series), min_periods=1).mean()
        else:
            series = series.ewm(self.N, min_periods=1).mean()
        return series


#################### Pair-Wise Rolling ####################
class PairRolling(OpsFeature):
    """Pair Rolling Operator
    
    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling output of two input features
    """
    def __init__(self, feature_left, feature_right, N, func):
        self.feature_left = feature_left
        self.feature_right = feature_right
        self.N = N
        self.func = func

    def __str__(self):
        return '{}_{}_{}_{}'.format(type(self).__name__.lower(), self.feature_left, self.feature_right, self.N)

    def _load_internal(self, instrument, start_time, end_time, freq):
        series_left = self.feature_left.load(instrument, start_time, end_time, freq)
        series_right = self.feature_right.load(instrument, start_time, end_time, freq)
        if self.N == 0:
            series = getattr(series_left.expanding(min_periods=1), self.func)(series_right)
        else:
            series = getattr(series_left.rolling(self.N, min_periods=1), self.func)(series_right)
        return series
    
    
class Corr(PairRolling):
    """Rolling Correlation
    
    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling correlation of two input features
    """
    def __init__(self, feature_left, feature_right, N):
        super(Corr, self).__init__(feature_left, feature_right, N, 'corr')
    

class Cov(PairRolling):
    """Rolling Covariance
    
    Parameters
    ----------
    feature_left : Feature
        feature instance
    feature_right : Feature
        feature instance
    N : int
        rolling window size

    Returns
    ----------
    Feature
        a feature instance with rolling max of two input features
    """
    def __init__(self, feature_left, feature_right, N):
        super(Cov, self).__init__(feature_left, feature_right, N, 'cov')
